﻿import sys, ast, os, linecache
from collections import namedtuple
"""
    Parse a .py file to a list of namedtuple objects
  
    Python objects of interest are represented by one of the following namedtuples:
"""
PythonObj = namedtuple('PythonObj', ('name',    # just the name, e.g. myvar
                                    'qualname',  # relative dotted name within this module, e.g. MyClass.myvar
                                    'docstr',
                                    'level',
                                    'lineno'))
Var = namedtuple('Var', PythonObj._fields)
Module = namedtuple('Module', PythonObj._fields + ('path',))
Class = namedtuple('Class', PythonObj._fields + ('bases',))
ClassEnd = namedtuple('ClassEnd', Class._fields)
Function = namedtuple('Function', PythonObj._fields + ('args',))
InstanceMethod = namedtuple('InstanceMethod', Function._fields + ('classname',))
InstanceVar = namedtuple('InstanceVar', Var._fields)
Property = namedtuple('Property', InstanceMethod._fields)
Constructor = namedtuple('Constructor', InstanceMethod._fields)


def parsemodule(modulepath):
    """ returns a list of namedtuple objects representing the Python objects of interest """
    parser = ParseModule(modulepath)
    return parser.parse()

def getbases(ast_ClassDef_node):
    """ given a ClassDef node from an AST, return a list-of-strings giving the bases """
    
    basenames = []
    for b in ast_ClassDef_node.bases:
        attrs = []
        while isinstance(b, ast.Attribute):
            attrs.insert(0, b.attr)
            b = b.value
        
        attrs.insert(0, b.id)
        basename = '.'.join(attrs)
        basenames.append(basename)
    
    return basenames

def gettarget(ast_Assign_node):
    """ takes an ast.Assign node and returns a string describing the (possibly-dotted) target
        
        can handle:
            foo = thing
            foo.bar = thing

        can't handle:
            foo, bar = thing
            foo.bar, baz = thing
    """
    
    target = ''
    for t in ast_Assign_node.targets:
        
        if isinstance(t, ast.Name): # have foo = bar
            target = t.id
        elif isinstance(t, ast.Attribute): # have foo.whatever.baz = bar
            attrs = []
            while isinstance(t, ast.Attribute):
                attrs.insert(0, t.attr)
                t = t.value
            attrs.insert(0, t.id)
            target = '.'.join(attrs)
            
    return target
        
class ParseModule(ast.NodeVisitor):
    
    def __init__(self, modulepath):
        
        if not modulepath.endswith('.py'):
            raise TypeError('can only document .py files, not "%s" files' % os.path.splitext(modulepath)[1])
        
        #print 'Documenting "%s"' % modulepath
        
        # internal state:
        self.modulepath = modulepath
        self.pyclassname = None   # name for the python class we're currently in
        self.pyclassproperties = [] # list of names for python class @property methods
        self.self_var = None # name of the method's "self" argument
        self.level = None
        self.rel_list = [] # the dotted name to the current item, as a list
        self.items = [] #  will store the things we find
        
        # parse the file to an ast:
        with open(modulepath, 'r') as modf:
            self.mod_ast = ast.parse(''.join(modf))
    
    def parse(self):
        
        # traverse the AST
        self.visit(self.mod_ast)
        
        return self.items
        
    def generic_visit(self, node):
        """ Like ast.generic_visit() but sets node.parent for children and self.level
        
            Calls visit() on all children of the node.
            In turn is called by visit() if no explicit visitor function is found for a node.
            Note that child nodes of nodes that have a custom visitor method won't
            be visited unless the visitor calls generic_visit() or visits them itself.
        """
        
        self.level += 1
        if not isinstance(self.items[-1], Module):
            self.rel_list.append(self.items[-1].name)
        
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        item.parent = node
                        self.visit(item)
            elif isinstance(value, ast.AST):
                value.parent = node
                self.visit(value)
        
        self.level -= 1
        if self.rel_list:
            self.rel_list.pop()
    
    def next_sibling(self, node):
        """ Returns the next ast node at the same "level", or None if there isn't one """
        parent = node.parent
        this = node
        siblings = list(ast.iter_child_nodes(parent))
        next_idx = siblings.index(this) + 1
        if next_idx > (len(siblings) - 1):
            return None
        else:
            return siblings[next_idx]
            
    def makequalname(self, name):
        return '.'.join(self.rel_list + [name])    
    
    # -- following methods are ast Node visitors, for ast.NodeVisitor --
    
    # def visit_Name(self, node):
        # print 'Name:', node.id, node.ctx
        # print
    # def visit_Str(self, node):
        # print 'string:', node.s
    
    def visit_Assign(self, node):
    
        """ Document assigns to:
            a. properties of the "self" variable in instance methods, e.g. 
                self.foo = bar
            b. module-level variables, e.g.
                baz = whatever
            
            For the above assigns to be considered they must include a "variable docstring" in one of the following ways
            (given in order of preference):
                1. "#: " single-line comment on same line
                2. "#: " single-line comment on preceding line
                3. bare string literal as next non-whitespace item (ignoring intervening whitespace)
        """
        
        #print 'assign to', node.targets, '@', node.lineno
        
        # get the IDs of the targets - if they're either a single one or a tuple
        target = gettarget(node)
        #print 'target:', target, repr(self.self_var)
        
        if target != '': # i.e. it is something we can handle
        
            # try getting docstring from same-line OR previous line comment:
            for lineno in (node.lineno, node.lineno - 1):
                sourceline = linecache.getline(self.modulepath, lineno)
                docstr = sourceline.partition('#:')[2].strip()
                if docstr != "":
                    break
            
            # try getting docstring from following bare string literal:
            if docstr == '':
                next_node = self.next_sibling(node)
                if isinstance(next_node, ast.Expr) and isinstance(next_node.value, ast.Str):
                    docstr = next_node.value.s
            
            if self.self_var is not None and target.startswith(self.self_var + '.'): # have an instance variable
                name = target[len(self.self_var) + 1:]
                level = self.level - 1
                qualname = '.'.join(self.rel_list[:-1] + [name])
                self.items.append(InstanceVar(name, qualname, docstr, level, node.lineno))
                return
            else: # have a standard variable
                name = target
                level = self.level
                qualname = self.makequalname(name)
                self.items.append(Var(name, qualname, docstr, level, node.lineno))
                return
            
    def visit_Module(self, node):
        self.level = 0
        name = os.path.splitext(os.path.basename(self.modulepath))[0]
        self.items.append(Module(name, '', ast.get_docstring(node), self.level, 0, self.modulepath)) # doesn't have a line number
        self.generic_visit(node) # visit all the children
        
    def visit_FunctionDef(self, node):
        # generic for all types of function:
        name = node.name
        decoratornames = [n.id for n in node.decorator_list if isinstance(n, ast.Name)]
        args = [a.id for a in node.args.args]
        docstr = ast.get_docstring(node)
        
        if self.pyclassname:  # have a function in a class
            
            ## TODO: handle @staticmethods where the below is not true!
            # strip the "self" argument
            if not len(args) >= 1:
                raise TypeError('your code is broken: function "%s" inside class "%s" must take at least one arg' % (name, self.pyclassname))
            self_var, args = args[0], args[1:]
            self.self_var = self_var

            if 'property' in decoratornames:
                self.pyclassproperties.append(name)
                self.items.append(Property(name, self.makequalname(name), docstr, self.level, node.lineno, args, self.pyclassname))
            ## TODO: handle @classmethod in future...
            elif name == '__init__':
                self.items.append(Constructor(name, self.makequalname(name), docstr, self.level, node.lineno, args, self.pyclassname))
                self.generic_visit(node)
            elif name in self.pyclassproperties:
                # suppress property setter/deleter methods
                pass
            else: # instance method
                self.items.append(InstanceMethod(name, self.makequalname(name), docstr, self.level, node.lineno, args, self.pyclassname))
            
            self.self_var = None
        else: # normal function
            self.items.append(Function(name, self.makequalname(name), docstr, self.level, node.lineno, args))

        return
        
    def visit_ClassDef(self, node):
        
        # built a string listing the class's bases:
        bases = getbases(node)
        
        classobj = Class(node.name, self.makequalname(node.name), ast.get_docstring(node), self.level, node.lineno, bases)
        self.items.append(classobj)
        
        # visit its children:
        self.pyclassname = node.name
        self.generic_visit(node)
        self.pyclassname = None
        self.pyclassproperties = []
        
        # create a ClassEnd object
        self.items.append(ClassEnd(**classobj._asdict()))
        
if __name__ == '__main__':
    modpath = sys.argv[1]
    parser = ParseModule(modpath)
    tree = ast.dump(parser.mod_ast)
    print tree
    
    #for it in parsemodule(modpath):
    #    print '(%i)' % it.level, '  '* it.level, it.__class__.__name__, repr(it.name), repr(it.qualname), it.lineno
    #    print 